// MIT License
// Copyright (c) 2022 - 傅莘莘
// Source URL: https://github.com/zjhellofss/KuiperInfer
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#include "runtime/runtime_ir.hpp"
#include <deque>
#include <iostream>
#include <memory>
#include <utility>
#include <vector>
#include "layer/abstract/layer_factory.hpp"
#include "runtime/runtime_ir.hpp"
#include "utils/time/time_logging.hpp"

namespace kuiper_infer {
//构造函数 (RuntimeGraph::RuntimeGraph): 初始化 RuntimeGraph 实例，设置模型参数路径和二进制文件路径。
//在 RuntimeGraph 类的构造函数中使用的 std::move(param_path) 和 std::move(bin_path) 是 C++ 中的一种技术，称为“移动语义”。这里的 std::move 是标准库中的一个函数模板，用于将其参数转换为右值引用。在这个上下文中，它的作用是优化内存使用和提高性能。
//用法: 在构造函数中，param_path 和 bin_path 是通过值传递的，这意味着当对象被构造时，它们的值被复制到新对象中。使用 std::move，我们可以避免这种复制，而是直接将它们的资源“移动”到新对象中。std::move 实际上是将对象转换为右值引用，使得在构造函数体内，可以利用移动构造器或移动赋值操作符接管这些值的资源。
RuntimeGraph::RuntimeGraph(std::string param_path, std::string bin_path)
    : param_path_(std::move(param_path)), bin_path_(std::move(bin_path)) {}

void RuntimeGraph::set_bin_path(const std::string& bin_path) { this->bin_path_ = bin_path; }

void RuntimeGraph::set_param_path(const std::string& param_path) { this->param_path_ = param_path; }

const std::string& RuntimeGraph::param_path() const { return this->param_path_; }

const std::string& RuntimeGraph::bin_path() const { return this->bin_path_; }
//Init 方法: 初始化运行时图，加载模型和构建图的操作算子。
bool RuntimeGraph::Init() {
  //if (this->bin_path_.empty() || this->param_path_.empty()) { - 检查二进制文件路径和参数文件路径是否为空。
  if (this->bin_path_.empty() || this->param_path_.empty()) {
    LOG(ERROR) << "The bin path or param path is empty";
    return false;
  }
  //创建一个新的 pnnx::Graph 对象，并使用智能指针进行管理。
  this->graph_ = std::make_unique<pnnx::Graph>();
  //调用 load 方法加载模型的参数和二进制文件。
  int32_t load_result = this->graph_->load(param_path_, bin_path_);
  if (load_result != 0) {
    LOG(ERROR) << "Can not find the param path or bin path: " << param_path_ << " " << bin_path_;
    return false;
  }
  //从图对象中获取操作符列表。
  std::vector<pnnx::Operator*> operators = this->graph_->ops;
  if (operators.empty()) {
    LOG(ERROR) << "Can not read the layers' define";
    return false;
  }
  // 清空当前对象的算子列表。
  operators_.clear();
  for (const pnnx::Operator* op : operators) {
    if (!op) {
      LOG(ERROR) << "Meet the empty node";
      continue;
    } else {
      //为每个非空操作符创建一个 RuntimeOperator 对象。
      std::shared_ptr<RuntimeOperator> runtime_operator = std::make_shared<RuntimeOperator>();
      // 初始化算子的名称
      runtime_operator->name = op->name;
      runtime_operator->type = op->type;

      // 初始化算子中的input
      InitGraphOperatorsInput(op->inputs, runtime_operator);

      // 记录输出operand中的名称
      InitGraphOperatorsOutput(op->outputs, runtime_operator);

      // 初始化算子中的attribute(权重)
      InitGraphAttrs(op->attrs, runtime_operator);

      // 初始化算子中的parameter
      InitGraphParams(op->params, runtime_operator);
      // 将初始化的算子添加到算子列表中。
      this->operators_.push_back(runtime_operator);
    }
  }

  graph_state_ = GraphState::NeedBuild;
  return true;
}
//Build 方法: 构建运行时图，包括创建节点关系、节点拓扑排序、初始化节点的输入输出空间。
void RuntimeGraph::Build() {
  //if (graph_state_ == GraphState::Complete) { - 检查当前图的状态是否已完成（GraphState::Complete），如果是，则不需要再次构建。
  if (graph_state_ == GraphState::Complete) {
    LOG(INFO) << "Model has been built already!";
    return;
  }
  //if (graph_state_ == GraphState::NeedInit) { - 如果图处于需要初始化的状态（GraphState::NeedInit）。
  if (graph_state_ == GraphState::NeedInit) {
    bool init_graph = Init();
    LOG_IF(FATAL, !init_graph || graph_state_ == GraphState::NeedInit) << "Init graph failed!";
  }
  //CHECK(graph_state_ >= GraphState::NeedBuild) - 确保图的状态至少为 NeedBuild。
  CHECK(graph_state_ >= GraphState::NeedBuild)
      << "Graph status error, current state is " << int32_t(graph_state_);
  LOG_IF(FATAL, this->operators_.empty()) << "Graph operators is empty, may be no init";

  // 构建节点关系
  CreateNodeRelation();

  // 节点拓扑排序
  ReverseTopoSort();

  // 初始化节点的输入和输出空间
  RuntimeOperatorUtils<float>::InitOperatorInput(operators_);
  RuntimeOperatorUtils<float>::InitOperatorOutput(graph_->ops, operators_);
  //graph_state_ = GraphState::Complete; - 将图的状态设置为完成。
  // GPT
  // 您提出了一个很好的问题。确实，在大多数情况下，构建（Build）神经网络模型之后，我们仍然需要保留这个模型以进行前向传播（Forward）。如果 Build 函数中有代码重置了图对象（graph_），这可能出于以下几个原因：

  // 临时对象：可能 graph_ 是一个用于构建或初始化过程中的临时对象。一旦构建过程完成，所需的信息可能已经被转移或复制到其他成员变量中，因此不再需要保留原始的 graph_ 对象。

  // 状态管理：在某些设计中，graph_ 对象的生命周期可能仅限于构建过程。构建完成后，其状态和数据可能被转换或转移至其他专门用于执行前向传播的对象或结构中。

  // 重新构建的需求：在某些情况下，每次执行前向传播之前可能需要根据不同的参数或配置重新构建图对象。因此，在每次 Build 调用之后释放 graph_ 对象，可以确保下一次构建时不会受到旧状态的影响。

  // 代码逻辑错误：也有可能这是代码中的逻辑错误或设计上的疏忽。如果在构建后仍然需要 graph_ 进行前向传播，那么重置 graph_ 对象可能是不合适的。

  // 代码上下文：没有完整的代码上下文，有时难以准确判断为什么会有特定的逻辑。可能还有其他上下文相关的原因导致了这种设计决策。

  graph_state_ = GraphState::Complete;
  if (graph_ != nullptr) {
    graph_.reset();
    graph_ = nullptr;
  }
}
//Forward 方法: 执行模型的前向传播，处理每个算子的前向计算。
void RuntimeGraph::Forward(bool debug) {
  // 检查当前的执行图是否已经初始化完毕
  if (graph_state_ < GraphState::Complete) {
    LOG(FATAL) << "Graph need be build!"
               << ", current state is " << int32_t(graph_state_);
  }

  if (debug) {
    // 初始化层时间状态收集器。
    utils::LayerTimeStatesSingleton::LayerTimeStatesCollectorInit();
  }
  //- 定义一个 lambda 函数，用于执行层的前向传播。
  auto forward_layer = [&](const std::shared_ptr<Layer<float>>& layer, const std::string& op_name,
                           const std::string& op_type) {
    //定义状态码。
    StatusCode status;
    if (debug) {
      {
        utils::LayerTimeLogging layer_time_logging(op_name, op_type);
        status = layer->Forward();
      }
    } else {
      //执行层的前向传播。
      status = layer->Forward();
    }
    return status;
  };
  // 遍历所有算子。
  for (const auto& current_op : operators_) {
    //初始化算子的前向标志为假
    current_op->has_forward = false;
    CHECK_GT(current_op->forward_index, 0);
    // 如果算子是输入或输出算子，则跳过执行。
    if (is_input_op(current_op->name) || is_output_op(current_op->name)) {
      current_op->has_forward = true;
      continue;
    }

    CHECK(current_op->layer != nullptr)
        << "The layer corresponding to the op " << current_op->name
        << " is empty, indicating that it may not have been created.";
    // 获取算子对应的层。
    std::shared_ptr<Layer<float>> layer = current_op->layer;
    //- 执行层的前向传播。
    StatusCode status = forward_layer(layer, current_op->name, current_op->type);
    CHECK(status == StatusCode::kSuccess)
        << layer->layer_name() << " layer forward failed, error code: " << int32_t(status);

    current_op->has_forward = true;
    // 将算子的输出传播到下一层算子。
    PropagateLayerOutputs(current_op, current_op->output_operands->datas);
  }

  if (debug) {
    utils::LayerTimeLogging::SummaryLogging();
  }

  for (const auto& op : operators_) {
    LOG_IF(FATAL, !op->has_forward) << "The operator: " << op->name << " has not been forward yet!";
  }
}

//CreateLayer 方法: 根据运行时操作创建相应的层。
template <typename T>
std::shared_ptr<Layer<T>> RuntimeGraph::CreateLayer(
    const std::shared_ptr<RuntimeOperatorBase<T>>& op) {
  LOG_IF(FATAL, !op) << "Operator is empty!";
  auto layer = LayerRegisterer::CreateLayer(op);
  LOG_IF(FATAL, !layer) << "Layer init failed " << op->type;
  return layer;
}

//InitGraphOperatorsInput, InitGraphOperatorsOutput, InitGraphParams, InitGraphAttrs 方法: 初始化算子的输入、输出、参数和属性。
template <typename T>
void RuntimeGraph::InitGraphOperatorsInput(
    const std::vector<pnnx::Operand*>& inputs,
    const std::shared_ptr<RuntimeOperatorBase<T>>& runtime_operator) {
  if (inputs.empty()) {
    return;
  }
  CHECK(runtime_operator != nullptr) << "The runtime operator is null pointer";
  for (const pnnx::Operand* input : inputs) {
    if (!input) {
      continue;
    }
    const pnnx::Operator* producer = input->producer;
    std::shared_ptr<RuntimeOperand> runtime_operand = std::make_shared<RuntimeOperand>();

    runtime_operand->name = producer->name;
    for (int32_t dim : input->shape) {
      runtime_operand->shapes.push_back(dim);
    }
    CHECK(!runtime_operand->shapes.empty());

    switch (input->type) {
      case 1: {
        runtime_operand->type = RuntimeDataType::kTypeFloat32;
        break;
      }
      case 0: {
        runtime_operand->type = RuntimeDataType::kTypeUnknown;
        break;
      }
      default: {
        LOG(FATAL) << "Unknown input operand type: " << input->type;
      }
    }

    runtime_operator->input_operands.insert({producer->name, runtime_operand});
    runtime_operator->input_operands_seq.push_back(runtime_operand);
  }
}

//InitGraphOperatorsInput, InitGraphOperatorsOutput, InitGraphParams, InitGraphAttrs 方法: 初始化算子的输入、输出、参数和属性。
template <typename T>
void RuntimeGraph::InitGraphOperatorsOutput(
    const std::vector<pnnx::Operand*>& outputs,
    const std::shared_ptr<RuntimeOperatorBase<T>>& runtime_operator) {
  if (outputs.empty()) {
    return;
  }
  CHECK(runtime_operator != nullptr) << "The runtime operator is null pointer";
  for (const pnnx::Operand* output : outputs) {
    if (!output) {
      continue;
    }
    const auto& consumers = output->consumers;
    for (const auto& c : consumers) {
      runtime_operator->output_names.push_back(c->name);
    }
  }
}
//InitGraphOperatorsInput, InitGraphOperatorsOutput, InitGraphParams, InitGraphAttrs 方法: 初始化算子的输入、输出、参数和属性。
template <typename T>
void RuntimeGraph::InitGraphParams(
    const std::map<std::string, pnnx::Parameter>& params,
    const std::shared_ptr<RuntimeOperatorBase<T>>& runtime_operator) {
  if (params.empty()) {
    return;
  }
  CHECK(runtime_operator != nullptr) << "The runtime operator is null pointer";
  for (const auto& [name, parameter] : params) {
    const int32_t type = parameter.type;
    switch (type) {
      case int32_t(RuntimeParameterType::kParameterUnknown): {
        std::shared_ptr<RuntimeParameter> runtime_parameter = std::make_shared<RuntimeParameter>();
        runtime_operator->params.insert({name, runtime_parameter});
        break;
      }

      case int32_t(RuntimeParameterType::kParameterBool): {
        std::shared_ptr<RuntimeParameterBool> runtime_parameter =
            std::make_shared<RuntimeParameterBool>(parameter.b);
        runtime_operator->params.insert({name, runtime_parameter});
        break;
      }

      case int32_t(RuntimeParameterType::kParameterInt): {
        std::shared_ptr<RuntimeParameterInt> runtime_parameter =
            std::make_shared<RuntimeParameterInt>(parameter.i);
        runtime_operator->params.insert({name, runtime_parameter});
        break;
      }

      case int32_t(RuntimeParameterType::kParameterFloat): {
        std::shared_ptr<RuntimeParameterFloat> runtime_parameter =
            std::make_shared<RuntimeParameterFloat>(parameter.f);
        runtime_operator->params.insert({name, runtime_parameter});
        break;
      }

      case int32_t(RuntimeParameterType::kParameterString): {
        std::shared_ptr<RuntimeParameterString> runtime_parameter =
            std::make_shared<RuntimeParameterString>(parameter.s);
        runtime_operator->params.insert({name, runtime_parameter});
        break;
      }

      case int32_t(RuntimeParameterType::kParameterIntArray): {
        std::shared_ptr<RuntimeParameterIntArray> runtime_parameter =
            std::make_shared<RuntimeParameterIntArray>(parameter.ai);
        runtime_operator->params.insert({name, runtime_parameter});
        break;
      }

      case int32_t(RuntimeParameterType::kParameterFloatArray): {
        std::shared_ptr<RuntimeParameterFloatArray> runtime_parameter =
            std::make_shared<RuntimeParameterFloatArray>(parameter.af);
        runtime_operator->params.insert({name, runtime_parameter});
        break;
      }
      case int32_t(RuntimeParameterType::kParameterStringArray): {
        std::shared_ptr<RuntimeParameterStringArray> runtime_parameter =
            std::make_shared<RuntimeParameterStringArray>(parameter.as);
        runtime_operator->params.insert({name, runtime_parameter});
        break;
      }
      default: {
        LOG(FATAL) << "Unknown parameter type: " << type;
      }
    }
  }
}
//InitGraphOperatorsInput, InitGraphOperatorsOutput, InitGraphParams, InitGraphAttrs 方法: 初始化算子的输入、输出、参数和属性。
template <typename T>
void RuntimeGraph::InitGraphAttrs(const std::map<std::string, pnnx::Attribute>& attrs,
                                  const std::shared_ptr<RuntimeOperatorBase<T>>& runtime_operator) {
  if (attrs.empty()) {
    return;
  }
  CHECK(runtime_operator != nullptr) << "The runtime operator is null pointer";
  for (const auto& [name, attr] : attrs) {
    switch (attr.type) {
      case 1: {
        std::shared_ptr<RuntimeAttribute> runtime_attribute = std::make_shared<RuntimeAttribute>(
            attr.shape, RuntimeDataType::kTypeFloat32, attr.data);
        runtime_operator->attribute.insert({name, runtime_attribute});
        break;
      }
      default: {
        LOG(FATAL) << "Unknown attribute type: " << attr.type;
      }
    }
  }
}
//PropagateLayerOutputs 方法: 传播层的输出到下一个操作的输入。
template <typename T>
void RuntimeGraph::PropagateLayerOutputs(
    const std::shared_ptr<RuntimeOperatorBase<T>>& current_op,
    const std::vector<std::shared_ptr<Tensor<T>>>& layer_output_datas) {
  // For each next operator of current operator
  //遍历当前算子的所有后继算子。output_operators 保存了与当前算子直接连接的下一层算子。
  for (const auto& [_, output_op] : current_op->output_operators) {
    // Get next op's input operands corresponding to current op's output
    // 获取后继算子的输入操作数列表。
    const auto& next_input_operands = output_op->input_operands;
    //在后继算子的输入操作数中查找当前算子的输出。
    const auto& next_input_op_iter = next_input_operands.find(current_op->name);
    // 如果找到了对应的输入操作数。
    if (next_input_op_iter != next_input_operands.end()) {
      // Get input data spaces for those operands
      //获取后继算子的输入数据缓冲区。
      std::vector<stensor<T>>& next_input_datas = next_input_op_iter->second->datas;
      // Copy current op output data to next op input data
      //- 遍历当前算子的输出数据。
      for (uint32_t i = 0; i < next_input_datas.size(); ++i) {
        //获取当前算子的第 i 个输出数据。
        const stensor<T>& layer_output_data = layer_output_datas.at(i);
        //检查后继算子的输入数据缓冲区是否已经初始化。
        if (next_input_datas.at(i) != nullptr) {
          //确保当前算子输出的数据形状与后继算子的输入数据形状匹配。
          CHECK(next_input_datas.at(i)->shapes() == layer_output_data->shapes());
        }
        //将当前算子的输出数据复制到后继算子的对应输入数据缓冲区中。
        next_input_datas.at(i) = layer_output_data;
      }
    }
  }
}
//ReverseTopoSort 和 ReverseTopoSortInternal 方法: 对图中的节点进行反向拓扑排序。
//RuntimeGraph::ReverseTopoSort 函数的主要目的是对神经网络模型中的算子进行反向拓扑排序。这是为了确定算子执行的顺序，确保在执行当前算子之前，所有依赖的输入算子都已经被执行
void RuntimeGraph::ReverseTopoSort() {
  // 构建拓扑顺序,遍历模型中的所有算子。
  for (const auto& op : operators_) {
    // 根据输入节点构建拓扑排序, 检查算子是否非空且还未进行拓扑排序。
    if (op != nullptr && !op->has_forward) {
      //初始化当前算子的前向索引。
      int32_t current_forward_idx = 0;
      //- 对当前算子及其依赖的算子进行递归的反向拓扑排序。
      this->ReverseTopoSortInternal(op, current_forward_idx);
    }
  }

  // 根据拓扑顺序调整算子的执行顺序,使用 lambda 表达式对所有算子根据它们的前向索引进行排序。
  // 确保执行顺序是按照依赖的反向顺序进行的。
  std::sort(operators_.begin(), operators_.end(), [](const auto& op1, const auto& op2) {
    return op1->forward_index > op2->forward_index;
  });
  //- 初始化前向索引。
  int32_t forward_index = 1;
  //- 再次遍历所有算子。
  for (const auto& op : operators_) {
    // 为每个算子分配一个新的前向索引。
    op->forward_index = forward_index;
    //- 增加前向索引。
    forward_index += 1;
  }
}
//ReverseTopoSort 和 ReverseTopoSortInternal 方法: 对图中的节点进行反向拓扑排序。
template <typename T>
void RuntimeGraph::ReverseTopoSortInternal(const std::shared_ptr<RuntimeOperatorBase<T>>& root_op,
                                           int32_t& current_forward_idx) {
  if (!root_op) {
    LOG(INFO) << "Current operator is nullptr";
    return;
  }
  if (root_op->input_operands.empty() && !root_op->has_forward) {
    this->input_ops_.push_back(root_op);
  }
  if (root_op->output_names.empty() && !root_op->has_forward) {
    this->output_ops_.push_back(root_op);
  }

  root_op->has_forward = true;
  const auto& next_ops = root_op->output_operators;
  for (const auto& [_, op] : next_ops) {
    if (op != nullptr && !op->has_forward) {
      this->ReverseTopoSortInternal(op, current_forward_idx);
    }
  }

  for (const auto& [_, op] : next_ops) {
    CHECK_EQ(op->has_forward, true);
  }
  root_op->forward_index = current_forward_idx;
  current_forward_idx += 1;
}
//RuntimeGraph::CreateNodeRelation 函数的主要目的是构建神经网络模型中各节点（算子）之间的关系。下面是该函数的逐行解释：
void RuntimeGraph::CreateNodeRelation() {
  // 构建图关系， 遍历 operators_ 中的每个算子（current_op）。
  for (const auto& current_op : this->operators_) {
    // 获取当前节点的所有后继节点的names，遍历根据next_op_name从operators_maps_中插入所需要的节点
    const std::vector<std::string>& output_names = current_op->output_names;
    //遍历每个输出节点的名称。
    for (const auto& kOutputName : output_names) {
      //再次遍历所有算子，寻找与输出节点名称相匹配的算子。
      for (const auto& output_op : this->operators_) {
        //如果找到名称匹配且不是当前算子本身的算子，则将其作为输出节点。
        if (output_op != current_op && output_op->name == kOutputName) {
          //将找到的算子作为当前算子的输出算子插入到 output_operators 映射中。
          current_op->output_operators.insert({kOutputName, output_op});
        }
      }
    }
    // 除了输入和输出节点，都创建layer
    if (current_op->type != "pnnx.Input" && current_op->type != "pnnx.Output") {
      //- 调用 CreateLayer 创建对应的层。
      std::shared_ptr<Layer<float>> layer = RuntimeGraph::CreateLayer(current_op);
      if (layer) {
        current_op->layer = layer;
        layer->set_runtime_operator(current_op);
      } else {
        LOG(FATAL) << "Layer " << current_op->name << " create failed!";
      }
    }
  }
}

RuntimeGraph::GraphState RuntimeGraph::graph_state() const { return this->graph_state_; }
//在这个过程中，函数首先确保整个模型已经构建完毕，然后根据输入的名称找到对应的输入算子，并将输入数据传递给这个算子。这是神经网络前向传播的准备步骤之一，确保模型可以接收并处理输入数据。
void RuntimeGraph::set_inputs(const std::string& input_name, const std::vector<sftensor>& inputs) {
  CHECK(this->graph_state_ == GraphState::Complete);
  std::shared_ptr<RuntimeOperator> input_op;
  for (auto op : this->input_ops_) {
    if (op->name == input_name) {
      input_op = op;
      break;
    }
  }
  CHECK(input_op != nullptr) << "Can not find the input operator: " << input_name;
  PropagateLayerOutputs(input_op, inputs);
}

std::vector<sftensor> RuntimeGraph::get_outputs(const std::string& output_name) const {
  CHECK(this->graph_state_ == GraphState::Complete);
  std::shared_ptr<RuntimeOperator> output_op;
  for (auto op : this->output_ops_) {
    if (op->name == output_name) {
      output_op = op;
    }
  }

  CHECK(output_op != nullptr) << "Can not find the output operator: " << output_name;
  std::vector<sftensor> outputs;
  for (const auto& input_operand : output_op->input_operands_seq) {
    std::copy(input_operand->datas.begin(), input_operand->datas.end(),
              std::back_inserter(outputs));
  }
  return outputs;
}

bool RuntimeGraph::is_input_op(const std::string& op_name) const {
  for (auto op : this->input_ops_) {
    CHECK(op != nullptr);
    if (op->name == op_name) {
      return true;
    }
  }
  return false;
}

bool RuntimeGraph::is_output_op(const std::string& op_name) const {
  for (auto op : this->output_ops_) {
    CHECK(op != nullptr);
    if (op->name == op_name) {
      return true;
    }
  }
  return false;
}

}  // namespace kuiper_infer
